# ---------------------安装所需的pip包--------------------------#
# pip install cnocr -i https://pypi.doubanio.com/simple        #
# pip install PyMuPDF                                          #
# --------------------------------------------------------------#

# !/usr/bin/python
# -*- coding: utf-8 -*-

import ctypes
import os
import shutil
import warnings

import fitz
from cnocr import CnOcr
from cnocr.utils import read_img


# ---把pdf每页都转换成一张图片---#
def pdf_to_img(input_pdf_name, output_file):
    pdf = fitz.open(input_pdf_name + '.pdf')
    page = pdf.page_count  # 读取pdf的页数
    for i in range(0, page):
        page_fn = pdf.load_page(i)  # 读取第i页
        pix = page_fn.get_pixmap()
        pix.save(os.path.join(output_file, str(i) + '.png'))  # 从0开始按序号另存页数为图片 v5.0修改：用函数链接输出路径，避免不同系统的分隔符差异带来的报错


# ---识别图片中的文字并且return---#
def img_ocr(input_pic):
    ocr = CnOcr()
    img_fp = input_pic  # 输入图片路径
    img = read_img(img_fp)  # 读取图片
    res = ocr.ocr(img)  # res是list（列表）
    len_res = len(res)  # 文本的行数
    text = ""
    for i in range(0, len_res):
        line_tuple = res[i]  # line_tuple是tuple（元组）
        len_line_tuple = len(line_tuple[0])
        s = ""
        single_line = s.join(line_tuple[0])
        text += single_line
    return text  # text为识别出来的文本，此处return可以给下面写入txt


# ---将识别出来的文字写入txt---#
def write_to_txt(input_pdf_name, text):
    with open(input_pdf_name + '.txt', 'a', encoding='utf-8') as file:  # 用a模式是续写而不是w(重新写入),编码为utf-8
        file.write(text)
        file.write('\n')  # 写完一张图片里面的内容就换行


# ---获取当前路径下所有pdf文件---#
def get_pdf_list():
    pdf_list = [x for x in os.listdir('.') if os.path.isfile(x) and os.path.splitext(x)[1] == '.pdf']
    return pdf_list


warnings.filterwarnings("ignore")  # 忽略waring输出（第一次运行的时候会有warning）
pdf_list = get_pdf_list()  # 获取当前路径下所有pdf的文件，返回值是list（列表） ##剩下代码均归入pdf_list的循环中（v2.0增加此功能）
for i in range(0, len(pdf_list)):
    name = os.path.splitext(pdf_list[i])[0]
    input_pdf_name = os.path.join(os.path.curdir, name)  # v5.0修改：用函数链接输出路径，避免不同系统的分隔符差异带来的报错
    output_file = os.path.join(os.path.curdir, 'pdf_2_image')  # v5.0修改：用函数链接输出路径，避免不同系统的分隔符差异带来的报错
    input_file = output_file  # pdf_to_img的output_file == img_ocr的input_pic
    os.mkdir(output_file)  # 生成临时文件夹存放pdf另存出来的img
    pdf_to_img(input_pdf_name, output_file)
    pic = os.listdir(input_file)  # 遍历临时文件夹下的图片,pic是list（列表）
    for pic_fn in pic:
        text = img_ocr(os.path.join(input_file, pic_fn))  # v5.0修改：用函数链接输出路径，避免不同系统的分隔符差异带来的报错
        write_to_txt(input_pdf_name, text)
    shutil.rmtree(output_file)  # 删除生成的临时文件夹和文件
ctypes.windll.user32.MessageBoxW(0, u'已完成PDF转化', u'已完成', 0)  # 弹出对话框提示完成任务
